/*
 * To change this license header, choose License Headers in Project Properties.
 * To change this template file, choose Tools | Templates
 * and open the template in the editor.
 */

package be.fedict.barth.signer;

import be.fedict.commons.eid.jca.BeIDProvider;
import java.io.IOException;
import java.io.InputStream;
import java.io.OutputStream;
import java.net.URL;
import java.security.KeyStore;
import java.security.KeyStoreException;
import java.security.MessageDigest;
import java.security.NoSuchAlgorithmException;
import java.security.NoSuchProviderException;
import java.security.PrivateKey;
import java.security.Security;
import java.security.UnrecoverableEntryException;
import java.security.cert.CRLException;
import java.security.cert.CertificateEncodingException;
import java.security.cert.CertificateException;
import java.security.cert.X509CRL;
import java.security.cert.X509Certificate;
import java.util.ArrayList;
import java.util.Arrays;
import java.util.logging.Level;
import java.util.logging.Logger;
import javax.naming.InvalidNameException;
import org.apache.commons.io.IOUtils;
import org.apache.pdfbox.cos.COSArray;
import org.apache.pdfbox.cos.COSDictionary;
import org.apache.pdfbox.cos.COSName;
import org.apache.pdfbox.cos.COSStream;
import org.apache.pdfbox.exceptions.SignatureException;
import org.apache.pdfbox.pdmodel.PDDocument;
import org.apache.pdfbox.pdmodel.interactive.digitalsignature.SignatureInterface;
import org.bouncycastle.asn1.ASN1EncodableVector;
import org.bouncycastle.asn1.DERSet;
import org.bouncycastle.asn1.cms.Attribute;
import org.bouncycastle.asn1.cms.AttributeTable;
import org.bouncycastle.asn1.ess.ESSCertIDv2;
import org.bouncycastle.asn1.ess.SigningCertificateV2;
import org.bouncycastle.asn1.nist.NISTObjectIdentifiers;
import org.bouncycastle.asn1.ocsp.OCSPObjectIdentifiers;
import org.bouncycastle.asn1.ocsp.OCSPResponse;
import org.bouncycastle.asn1.pkcs.PKCSObjectIdentifiers;
import org.bouncycastle.asn1.x509.AlgorithmIdentifier;
import org.bouncycastle.cert.jcajce.JcaCRLStore;
import org.bouncycastle.cms.CMSProcessableByteArray;
import org.bouncycastle.cms.CMSSignedDataGenerator;
import org.bouncycastle.cert.jcajce.JcaCertStore;
import org.bouncycastle.cert.ocsp.OCSPException;
import org.bouncycastle.cms.CMSAttributeTableGenerator;
import org.bouncycastle.cms.CMSException;
import org.bouncycastle.cms.CMSSignedData;
import org.bouncycastle.cms.DefaultSignedAttributeTableGenerator;
import org.bouncycastle.cms.SignerInfoGenerator;
import org.bouncycastle.cms.SignerInformation;
import org.bouncycastle.cms.SignerInformationStore;
import org.bouncycastle.cms.jcajce.JcaSignerInfoGeneratorBuilder;
import org.bouncycastle.jce.provider.BouncyCastleProvider;
import org.bouncycastle.operator.ContentSigner;
import org.bouncycastle.operator.DigestCalculatorProvider;
import org.bouncycastle.operator.OperatorCreationException;
import org.bouncycastle.operator.bc.BcDigestCalculatorProvider;
import org.bouncycastle.operator.jcajce.JcaContentSignerBuilder;
import org.bouncycastle.tsp.TSPException;
import org.bouncycastle.tsp.TimeStampResponse;
import org.bouncycastle.tsp.TimeStampToken;
import org.bouncycastle.util.Store;

/**
 *
 * @author hanssens_bart
 */
public class PadesSigner implements SignatureInterface {
    private PrivateKey privateKey;
    private X509Certificate[] chain;
    private X509Certificate certificate;
    private OCSPResponse ocsp;
    private TimeStampResponse timestamp;
    private X509CRL[] crls;
    
       
    public String getKeys() throws KeyStoreException, IOException, NoSuchAlgorithmException, UnrecoverableEntryException, CertificateException, InvalidNameException {
        KeyStore store = KeyStore.getInstance("BeID");
        store.load(null);
        KeyStore.PrivateKeyEntry privKeyEntry = (KeyStore.PrivateKeyEntry) 
                store.getEntry("Signature", null);
        privateKey = privKeyEntry.getPrivateKey();
        chain = (X509Certificate[]) privKeyEntry.getCertificateChain();
        certificate = (X509Certificate) privKeyEntry.getCertificate();

        return Util.getSignerName(certificate);
    }
    
    public void getRevocation() throws CertificateEncodingException, IOException, OCSPException, NoSuchAlgorithmException, NoSuchProviderException, OperatorCreationException, CertificateException, CRLException {
        ocsp = getOCSPRevocation();
        crls = getCRLs();
    }
    
    private X509CRL[] getCRLs() throws IOException, CertificateException, CRLException {
        CRLGetter crler = new CRLGetter();
        return crler.get(certificate, chain);
    }
    
    private CMSAttributeTableGenerator getSignedAttribGenerator() throws NoSuchAlgorithmException, CertificateEncodingException, NoSuchProviderException, IOException, OCSPException, OperatorCreationException  {
        ASN1EncodableVector attributes = new ASN1EncodableVector();
        attributes.add(getSigningCertAttribute(certificate));
        attributes.add(getOCSPAttribute());
        AttributeTable table = new AttributeTable(attributes);
        
        CMSAttributeTableGenerator generator = new DefaultSignedAttributeTableGenerator(table);

        return generator;
    }
    
    private OCSPResponse getOCSPRevocation() throws CertificateEncodingException, IOException, NoSuchAlgorithmException, NoSuchProviderException, OCSPException, OperatorCreationException  {
        OCSPChecker checker = new OCSPChecker();
        return checker.check(certificate, chain);
    }
    
    private Attribute getOCSPAttribute() throws CertificateEncodingException, IOException, OCSPException, NoSuchAlgorithmException, NoSuchProviderException, OperatorCreationException {
        return new Attribute(OCSPObjectIdentifiers.id_pkix_ocsp, new DERSet(this.ocsp));
    }
    
    private TimeStampResponse getTimestamp(byte[] hash) throws IOException, TSPException {
        TimeStamper stamper = new TimeStamper();
        return stamper.stamp(new URL("http://tsa.belgium.be/connect"), hash);
    }

    private Attribute getTimestampAttribute() {
        TimeStampToken token = this.timestamp.getTimeStampToken();
        DERSet der = new DERSet(token.toCMSSignedData().toASN1Structure());
   
        return new Attribute(PKCSObjectIdentifiers.id_aa_signatureTimeStampToken, der);
    }
    
    private CMSSignedData addTimeStamp(CMSSignedData sigData) throws NoSuchAlgorithmException, NoSuchProviderException, IOException, TSPException {
        ASN1EncodableVector attributes = new ASN1EncodableVector();
        ArrayList<SignerInformation> storeNew = new ArrayList<SignerInformation>();
        
        MessageDigest digest = MessageDigest.getInstance("SHA256", "BC");

        for (Object o : sigData.getSignerInfos().getSigners()) {
            SignerInformation siOld = (SignerInformation) o;
            
            byte[] hash = digest.digest(siOld.getSignature());
            this.timestamp = getTimestamp(hash);
            attributes.add(getTimestampAttribute());
            AttributeTable table = new AttributeTable(attributes);
            
            storeNew.add(SignerInformation.replaceUnsignedAttributes(siOld, table));            
        }     
        
        return CMSSignedData.replaceSigners(sigData, new SignerInformationStore(storeNew));
    }
    
    private Attribute getSigningCertAttribute(X509Certificate cert) throws NoSuchAlgorithmException, CertificateEncodingException, NoSuchProviderException {
        MessageDigest digest = MessageDigest.getInstance("SHA256", "BC");
        byte[] hash = digest.digest(cert.getEncoded());

        ESSCertIDv2 essCert = new ESSCertIDv2(
                new AlgorithmIdentifier(NISTObjectIdentifiers.id_sha256), hash);  
        SigningCertificateV2 scv2 = new SigningCertificateV2(new ESSCertIDv2[] { essCert });

        return new Attribute(PKCSObjectIdentifiers.id_aa_signingCertificateV2, new DERSet(scv2));
    }
    
    public void addDSS(PDDocument pdf) throws CertificateEncodingException, IOException, CRLException {
        COSArray certs = new COSArray();
        
        for (X509Certificate c: chain) {
            byte[] b = c.getEncoded();
            COSStream cos = pdf.getDocument().createCOSStream();
            OutputStream out = cos.createUnfilteredStream();
            out.write(b);
            out.flush();
            out.close();
            
            certs.add(cos.getCOSObject());
        }
  
        COSArray CRLS = new COSArray();
        
        for (X509CRL c: crls) {
            byte[] b = c.getEncoded();
            COSStream cos = pdf.getDocument().createCOSStream();
            OutputStream out = cos.createUnfilteredStream();
            out.write(b);
            out.flush();
            out.close();
            
            CRLS.add(cos.getCOSObject());
        }
        
        COSArray ocsps = new COSArray();
        COSStream cos = pdf.getDocument().createCOSStream();
        OutputStream out = cos.createFilteredStream();
        out.write(ocsp.getEncoded());
        out.flush();
        out.close();
        
        ocsps.add(cos.getCOSObject());
        
        COSDictionary dict = new COSDictionary();
        dict.setName(COSName.TYPE, "DSS"); 
        dict.setItem("Certs", certs);
        dict.setItem("OCSPs", ocsps);
        dict.setItem("CRLs", CRLS);
     
        COSDictionary ext = new COSDictionary();
        ext.setName("BaseVersion", "1.7");
        ext.setInt("ExtensionLevel", 5);
        ext.setDirect(true);
        
        COSDictionary adbe = new COSDictionary();
        adbe.setItem("ADBE", ext);
        adbe.setDirect(true);
        
        COSDictionary sig = pdf.getDocumentCatalog().getCOSDictionary();
        sig.setItem("DSS", dict);
        sig.setItem("Extensions", adbe);
    }
    
    public byte[] sign(InputStream stream) throws SignatureException, IOException {
        CMSSignedData sigData = null;
        
        try {
            CMSProcessableByteArray input = 
                    new CMSProcessableByteArray(IOUtils.toByteArray(stream));
            
            ContentSigner sha = new JcaContentSignerBuilder("SHA256WithRSA")
                    .setProvider("BeIDProvider").build(privateKey);
            DigestCalculatorProvider digester = new BcDigestCalculatorProvider();
            JcaSignerInfoGeneratorBuilder sigbuilder = 
                    new JcaSignerInfoGeneratorBuilder(digester);
            sigbuilder.setSignedAttributeGenerator(getSignedAttribGenerator());
 
            SignerInfoGenerator siginfo = sigbuilder.build(sha, certificate);            
            
            CMSSignedDataGenerator gen = new CMSSignedDataGenerator();
            Store certStore = new JcaCertStore(Arrays.asList(chain));
            Store crlStore = new JcaCRLStore(Arrays.asList(crls));
        
            gen.addSignerInfoGenerator(siginfo);
            gen.addCertificates(certStore);
            gen.addCRLs(crlStore);
            
            sigData = gen.generate(input);
            sigData = addTimeStamp(sigData);
            
        } catch (CertificateEncodingException ex) {
            throw new IOException(ex);
        } catch (CMSException ex) {
            throw new IOException(ex);
        } catch (OperatorCreationException ex) {
            throw new IOException(ex);
        } catch (NoSuchAlgorithmException ex) {
            throw new IOException(ex);
        } catch (NoSuchProviderException ex) {
            throw new IOException(ex);
        } catch (TSPException ex) {
            throw new IOException(ex);
        } catch (OCSPException ex) {
            throw new IOException(ex);
        } catch (CRLException ex) {
            Logger.getLogger(PadesSigner.class.getName()).log(Level.SEVERE, null, ex);
        }
        return sigData.getEncoded();
    }
    
    public PadesSigner() {
        Security.addProvider(new BeIDProvider());
        Security.addProvider(new BouncyCastleProvider());
    }
}
